import dataclasses
import datetime
import itertools
import logging
from typing import List, Dict, OrderedDict, DefaultDict, Optional
from absl import app
import collections
import os
import time

from coin.base.datetime_util import (to_datetime, to_timestamp_int)
from coin.proto.coin_market_query_pb2 import ProductTicker
from coin.proto.coin_strategy_pb2 import StrategyLog
from coin.support.proto_log.logic.util import (read_strat_log, StratInfo)

os.environ["TZ"] = "GMT"
time.tzset()


@dataclasses.dataclass
class ReserveChange:
  ts: int
  qty: float
  price: Optional[ProductTicker]


class OnLog(object):
  def __init__(self, hypo_reserve_info: OrderedDict[datetime.datetime, Dict[str, Optional[float]]]):
    assert len(hypo_reserve_info) > 0, 'hypo reserve info should not be empty'
    self._get_hypo_info(hypo_reserve_info)
    self._get_base_and_symbol()
    self._last_reserve: Dict[str, Optional[float]] = {}
    self._last_ticker: Dict[str, Optional[ProductTicker]] = {}
    self._reserve_changes: DefaultDict[str, List[ReserveChange]] = collections.defaultdict(list)

  def _get_hypo_info(self, hypo_reserve_info):
    self._hypo_reserve: DefaultDict[str, List[ReserveChange]] = collections.defaultdict(list)
    self._hypo_reserve_checkpoint: Dict[str, int] = {}
    for ts, reserve_info in hypo_reserve_info.items():
      for product_str, qty in reserve_info.items():
        self._hypo_reserve[product_str].append(ReserveChange(to_timestamp_int(ts), qty, None))
    for product_str, reserve_changes in self._hypo_reserve.items():
      assert len(reserve_changes) > 0, 'reserve changes should not be empty'
      self._hypo_reserve_checkpoint[product_str] = 0
      for reserve_change in reserve_changes:
        if reserve_change.qty is None:
          self._hypo_reserve_checkpoint[product_str] += 1

  def _get_base_and_symbol(self):
    symbol_set = {product_str.split(':')[2] for product_str in self._hypo_reserve.keys()}
    assert len(symbol_set) == 1, 'base should be the same'
    self._symbol = symbol_set.pop()
    self._base = self._symbol.split('-')[0]

  def _assign_product_ticker(self, ticker: ProductTicker, product_str: str):
    hypo_reserve_list = self._hypo_reserve[product_str]
    checkpoint = self._hypo_reserve_checkpoint[product_str]
    ticker_ts = ticker.exchange_timestamp
    if len(hypo_reserve_list) > checkpoint and hypo_reserve_list[checkpoint].price is None and abs(
        to_datetime(ticker_ts) -
        to_datetime(hypo_reserve_list[checkpoint].ts)) < datetime.timedelta(minutes=5):
      logging.info(f'[{product_str}] hypo reserve price assinged: {to_datetime(ticker_ts)}'
                   f' {ticker.last}@{hypo_reserve_list[checkpoint].qty}')
      hypo_reserve_list[checkpoint].price = ticker
      self._hypo_reserve_checkpoint[product_str] += 1
    reserve_change_list = self._reserve_changes[product_str]
    if len(reserve_change_list) > 0:
      reserve_change_list[-1].price = ticker

  def _handle_reserve(self, ts: int, pb: StrategyLog):
    for each_balance in pb.reserve.each_balance:
      if each_balance.currency == self._base:
        total = each_balance.total
        product_str = f"{pb.reserve.market_type}:{pb.reserve.exchange}:{self._symbol}"
        if self._last_reserve.get(product_str) is None or self._last_reserve[product_str] != total:
          logging.info(f'[{product_str}] reserve changed: {to_datetime(ts)}'
                       f' from {self._last_reserve.get(product_str, 0.0)} to {total}')
          reserve_change = ReserveChange(ts, total, self._last_ticker.get(product_str))
          self._reserve_changes[product_str].append(reserve_change)
        self._last_reserve[product_str] = total

  def _handle_ticker(self, ts: int, pb: StrategyLog):
    for each_ticker in pb.ticker.each_ticker:
      if each_ticker.symbol == self._symbol:
        product_str = f"{pb.ticker.market_type}:{pb.ticker.exchange}:{self._symbol}"
        each_ticker.exchange_timestamp = ts
        self._assign_product_ticker(each_ticker, product_str)
        self._last_ticker[product_str] = each_ticker

  def on_log(self, ts: int, log: bytes):
    pb = StrategyLog()
    pb.ParseFromString(log)
    if pb.type == StrategyLog.RESERVE:
      self._handle_reserve(ts, pb)
    elif pb.type == StrategyLog.TICKER:
      self._handle_ticker(ts, pb)

  def get_corp_offset_pnl(self) -> float:
    corp_offset_pnl = dict(zip(self._reserve_changes.keys(), itertools.repeat(0.0)))

    for product_str, reserve_changes in self._reserve_changes.items():
      hypo_reserves = self._hypo_reserve[product_str]
      idx = 0
      cur_hypo_reserve = hypo_reserves[idx]
      while cur_hypo_reserve.qty is None and idx + 1 < len(hypo_reserves):
        idx += 1
        cur_hypo_reserve = hypo_reserves[idx]
      if len(reserve_changes) == 0:
        raise ValueError(f'no reserve change for {product_str}')
      assert reserve_changes[0].price.exchange_timestamp >= cur_hypo_reserve.ts, (
          f'price timestamp {reserve_changes[0].price.exchange_timestamp} should be'
          f' greater than reserve timestamp {cur_hypo_reserve.ts}')
      assert reserve_changes[0].price is not None and cur_hypo_reserve.price is not None, (
          'price should not be None')
      # first offset pnl must be calculated between hypo_reserve price and reserve_change price
      corp_offset_pnl[product_str] += (reserve_changes[0].price.last - cur_hypo_reserve.price.last)\
        * (reserve_changes[0].qty - cur_hypo_reserve.qty)
      logging.info(
          f'[{product_str}] first offset pnl: ({reserve_changes[0].price.last} - '
          f'{cur_hypo_reserve.price.last}) * ({reserve_changes[0].qty} - {cur_hypo_reserve.qty}) = '
          f'{corp_offset_pnl[product_str]}')
      for x, y in itertools.pairwise(reserve_changes):
        assert x.price is not None and y.price is not None, 'price should not be None'
        if idx + 1 < len(hypo_reserves) and x.ts < hypo_reserves[
            idx + 1].ts and y.ts >= hypo_reserves[
                idx + 1].ts:  # hypo reserve changed between offset change
          corp_offset_pnl[product_str] += (hypo_reserves[idx + 1].price.last -
                                           x.price.last) * (x.qty - cur_hypo_reserve.qty)
          logging.info(
              f'[{product_str}] offset pnl: ({hypo_reserves[idx + 1].price.last} - '
              f'{x.price.last}) * ({x.qty} - {cur_hypo_reserve.qty}) = '
              f'{(hypo_reserves[idx + 1].price.last - x.price.last) * (x.qty - cur_hypo_reserve.qty)}'
          )
          idx += 1
          cur_hypo_reserve = hypo_reserves[idx]
          assert cur_hypo_reserve.price is not None, 'price should not be None'
          corp_offset_pnl[product_str] += (y.price.last - cur_hypo_reserve.price.last) * (
              y.qty - cur_hypo_reserve.qty)
          logging.info(
              f'[{product_str}] offset pnl: ({y.price.last} - '
              f'{cur_hypo_reserve.price.last}) * ({y.qty} - {cur_hypo_reserve.qty}) = '
              f'{(y.price.last - cur_hypo_reserve.price.last) * (y.qty - cur_hypo_reserve.qty)}')
        else:  # hypo reserve doesn't change between offset change
          corp_offset_pnl[product_str] += (y.price.last - x.price.last) * (y.qty -
                                                                           cur_hypo_reserve.qty)
          logging.info(f'[{product_str}] offset pnl: {y.price.last} - '
                       f'{x.price.last} * ({y.qty} - {cur_hypo_reserve.qty}) = '
                       f'{(y.price.last - x.price.last) * (y.qty - cur_hypo_reserve.qty)}')
    logging.info(f'corp offset pnl: {corp_offset_pnl}')
    return corp_offset_pnl


def main(_):
  strategy_name = 'vmm_ccx_delta_hedge'
  machine = 'strategy-21.ap-southeast-1'
  root_dir = '/remote/iosg/strat-1/buckets/log.raw.coin/live/strat_proto_log'
  start_time = datetime.datetime(2023, 1, 1)
  end_time = datetime.datetime(2023, 3, 23)
  pnl_diff_calc = OnLog(
      collections.OrderedDict([
          (datetime.datetime.strptime("2023-01-18 09:58:12", "%Y-%m-%d %H:%M:%S"), {
              "Spot:Phemex:CCX-USDT": 12_000_000.0,
              "Spot:Bitmart:CCX-USDT": None
          }),
          (datetime.datetime.strptime("2023-03-06 10:05:09", "%Y-%m-%d %H:%M:%S"), {
              "Spot:Phemex:CCX-USDT": 6_000_000.0,
              "Spot:Bitmart:CCX-USDT": 6_000_000.0
          }),
      ]))
  strat_info = StratInfo(strategy_name=strategy_name, machine=machine, trading_date=None)
  read_strat_log(strat_info=strat_info,
                 start_time=start_time,
                 end_time=end_time,
                 callback=pnl_diff_calc.on_log,
                 root_dir=root_dir)
  print(pnl_diff_calc.get_corp_offset_pnl())


if __name__ == '__main__':

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
